#include <iostream>
#include <fstream>
#include <string>
using namespace std;

#include "../config.h"

#include "../RLE/RLE.h"
#include "../BWT/BWT.h"
#include "../ModeToFront/MTF.h"
#include "../Huffman/Huffman.h"
#include "../Adler-32 checksum/Adler32.h"

#define SAVE_STEPS

void save(string fname, const unsigned char *data, unsigned int len);

int main(int argc, char **argv)
{
    if(argc != 3)
    {
        cout << "Usage: ./prc_seq file_to_compress.any       compressed_file.any.prc\n"
             << " or    ./prc_seq file_to_decompress.any.prc decompressed_file.any" << endl;
        return 1;
    }
    unsigned char *buffer = new unsigned char[Config::BLOCK_SIZE];
    ifstream fin(argv[1], ios_base::in | ios_base::binary);
    if(!fin.is_open()) { cerr << "File reading error!" << endl; return 1; }
    ofstream fout(argv[2], ios_base::out | ios_base::binary);
    if(!fout.is_open()) { cerr << "File writing error!" << endl; return 1; }
    // compression or decompression??
    bool decompress = false;
    if(strcmp(argv[1] + strlen(argv[1]) - 4, ".prc") == 0)
        decompress = true;
    //
    streamsize read;
    unsigned int bwt_dec_start, enc_len, dec_len, checksum;
    unsigned char *encoded = new unsigned char[Config::BLOCK_SIZE];
    unsigned char *decoded = new unsigned char[Config::BLOCK_SIZE];
    unsigned short *rle_ab_enc = new unsigned short[Config::BLOCK_SIZE];
    //
    // DECOMPRESSION
    if(decompress)
    {
        // check file type
        fin.read((char *)buffer, 3);
        if(buffer[0] != 'P' || buffer[1] != 'R' || buffer[2] != 'C')
        {
            cerr << "Invalid format!" << endl;
            return 1;
        }
        // decompression
        while(1)
        {
            fin.read((char *)&enc_len, sizeof(unsigned int));		// read a length of a compressed block
            fin.read((char *)buffer, enc_len+sizeof(unsigned int));	// read a compressed block of length `enc_len` and read the checksum for this block
            if((read = fin.gcount()) > 0)	// get the count of read bytes in the last read operation
            {
                enc_len = read - 4;
                if(!Adler32::Check(buffer, enc_len, *((unsigned int *)(buffer+enc_len))))	// read and check the checksum
                {
                    cerr << "File is corrupted. Invalid checksum!" << endl;
                    return 2;
                }
                // decompress one block
                Huffman::Decode(enc_len, buffer, rle_ab_enc, enc_len);
                RLE::Decode_AB(rle_ab_enc, enc_len, encoded, enc_len);
                bwt_dec_start = ((int *)encoded)[0]; enc_len -= 4;
                MTF::Decode(encoded+4, enc_len, decoded);
                BWT::Decode(enc_len, decoded, bwt_dec_start, encoded);
                RLE::Decode_MN5(encoded, enc_len, decoded, dec_len);
                // write out the decompressed data
                fout.write((char *)decoded, dec_len);
            }
            if(fin.eof()) break;
        }
    }
    //
    // COMPRESSION
    else
    {
        fout.write("PRC", 3);
        while(1)
        {
            fin.read((char *)buffer, Config::BLOCK_SIZE);
            if((read = fin.gcount()) > 0)
            {
                RLE::Encode_MN5(buffer, read, encoded, enc_len);
                save("step1.mn5", encoded, enc_len);
                BWT::Encode(enc_len, encoded, buffer, bwt_dec_start);
                save("step2.bwt", buffer, enc_len);
                MTF::Encode(buffer, enc_len, encoded+4);
                ((int *)encoded)[0] = bwt_dec_start; enc_len += 4;
                save("step3.mtf", encoded, enc_len);
                RLE::Encode_AB(encoded, enc_len, rle_ab_enc, enc_len);
                save("step4.ab", (unsigned char*)rle_ab_enc, enc_len*sizeof(short));
                Huffman::Encode(enc_len, rle_ab_enc, encoded, enc_len);
                save("step5.huf", encoded, enc_len);
                #ifdef SAVE_STEPS
                    break;
                #endif
                checksum = Adler32::Compute(encoded, enc_len);
                // write out
                fout.write((char *)&enc_len, sizeof(unsigned int));	// encoded block length (for faster decompression - especially in parallel)
                fout.write((char *)encoded, enc_len);				// encoded block
                fout.write((char *)&checksum, sizeof(unsigned int));// checksum of the encoded block
            }
            if(fin.eof()) break;
        }
    }
    fout.close();
    fin.close();
    //
    delete [] buffer;
    delete [] encoded;
    delete [] decoded;
    delete [] rle_ab_enc;
    //
    return 0;
}

void save(string fname, const unsigned char *data, unsigned int len)
{
    #ifdef SAVE_STEPS
        ofstream fout(fname.c_str(), ios_base::out | ios_base::binary);
        fout.write((const char *)data, len);
        fout.close();
    #endif
}